package com.gitee.xmhzzz.component.spring.cloud.component.dao.interceptor;

import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.NotEqualsTo;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.select.Join;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SubSelect;
import net.sf.jsqlparser.statement.update.Update;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.*;

/**
 * @ClassName SqlCheckInterceptor
 * @Description sql检测
 * @Author wzq
 * @Date 2023/9/24 21:44
 * @Version 1.0
 */
public class SqlCheckInterceptor extends AbstractSqlCheckJsqlParserSupport {
    private static final Logger log = LoggerFactory.getLogger(SqlCheckInterceptor.class);
    private static final String SQL_CHECK_PREFIX = "mybatis sql check {} ";
    private static final Set<String> CACHE_VALID_RESULT = new HashSet();

    public SqlCheckInterceptor() {
    }

    public void beforePrepare(String originalSql, Long checkLimit, int numRows, Long checkRows) {
        String md5Base64 = md5Base64(originalSql);
        if (!CACHE_VALID_RESULT.contains(md5Base64)) {
            if ((long)numRows >= checkRows) {
                log.warn("mybatis sql check 返回行数大于{}", numRows);
            }

            this.parserSingle(originalSql, checkLimit);
            CACHE_VALID_RESULT.add(md5Base64);
        }
    }

    protected void processSelect(Select select, int index, String sql, Long checkLimit) {
        PlainSelect plainSelect = (PlainSelect)select.getSelectBody();
        Expression where = plainSelect.getWhere();
        Optional.ofNullable(plainSelect.getLimit()).ifPresent((limit) -> {
            long rowCount = ((LongValue)limit.getRowCount(LongValue.class)).getValue();
            if (rowCount >= checkLimit) {
                log.warn("mybatis sql check {} ", "limit查询行数大于" + rowCount);
            }

        });
        if (where == null) {
            log.error("mybatis sql check {} ", "非法SQL，必须要有where条件");
        } else {
            Table table = (Table)plainSelect.getFromItem();
            List<Join> joins = plainSelect.getJoins();
            this.validWhere(where, table);
            this.validJoins(joins, table);
        }
    }

    protected void processUpdate(Update update, int index, String sql) {
        Expression where = update.getWhere();
        if (where == null) {
            log.error("mybatis sql check {} ", "非法SQL，必须要有where条件");
        } else {
            Table table = update.getTable();
            List<Join> joins = update.getJoins();
            this.validWhere(where, table);
            this.validJoins(joins, table);
        }
    }

    protected void processDelete(Delete delete, int index, String sql) {
        Expression where = delete.getWhere();
        if (where == null) {
            log.error("mybatis sql check {} ", "非法SQL，必须要有where条件");
        } else {
            Table table = delete.getTable();
            List<Join> joins = delete.getJoins();
            this.validWhere(where, table);
            this.validJoins(joins, table);
        }
    }

    private void validExpression(Expression expression) {
        if (expression instanceof OrExpression) {
            OrExpression orExpression = (OrExpression)expression;
            log.error("mybatis sql check {} ", "非法SQL，where条件中不能使用【or】关键字，错误or信息：" + orExpression.toString());
        } else if (expression instanceof NotEqualsTo) {
            NotEqualsTo notEqualsTo = (NotEqualsTo)expression;
            log.error("mybatis sql check {} ", "非法SQL，where条件中不能使用【!=】关键字，错误!=信息：" + notEqualsTo.toString());
        } else {
            SubSelect subSelect;
            if (expression instanceof BinaryExpression) {
                BinaryExpression binaryExpression = (BinaryExpression)expression;
                if (binaryExpression.getLeftExpression() instanceof Function) {
                    Function function = (Function)binaryExpression.getLeftExpression();
                    log.error("mybatis sql check {} ", "非法SQL，where条件中不能使用数据库函数，错误函数信息：" + function.toString());
                }

                if (binaryExpression.getRightExpression() instanceof SubSelect) {
                    subSelect = (SubSelect)binaryExpression.getRightExpression();
                    log.error("mybatis sql check {} ", "非法SQL，where条件中不能使用子查询，错误子查询SQL信息：" + subSelect.toString());
                }
            } else if (expression instanceof InExpression) {
                InExpression inExpression = (InExpression)expression;
                if (inExpression.getRightItemsList() instanceof SubSelect) {
                    subSelect = (SubSelect)inExpression.getRightItemsList();
                    log.error("mybatis sql check {} ", "非法SQL，where条件中不能使用子查询，错误子查询SQL信息：" + subSelect.toString());
                }
            }
        }

    }

    private void validJoins(List<Join> joins, Table table) {
        if (joins != null) {
            Iterator var3 = joins.iterator();

            while(var3.hasNext()) {
                Join join = (Join)var3.next();
                Table rightTable = (Table)join.getRightItem();
                Expression expression = join.getOnExpression();
                this.validWhere(expression, table, rightTable);
            }
        }

    }

    private void validWhere(Expression expression, Table table) {
        this.validWhere(expression, table, (Table)null);
    }

    private void validWhere(Expression expression, Table table, Table joinTable) {
        this.validExpression(expression);
        if (expression instanceof BinaryExpression) {
            Expression leftExpression = ((BinaryExpression)expression).getLeftExpression();
            this.validExpression(leftExpression);
            if (leftExpression instanceof BinaryExpression) {
                this.validWhere(leftExpression, table, joinTable);
            }

            Expression rightExpression = ((BinaryExpression)expression).getRightExpression();
            this.validExpression(rightExpression);
        }

    }

    public static String md5Base64(String str) {
        try {
            MessageDigest md5 = MessageDigest.getInstance("MD5");
            byte[] src = md5.digest(str.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(src);
        } catch (Exception var3) {
            log.error("加密失败", var3);
            return "";
        }
    }
}
